# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 结构化数据-在多个模式之间选择.py
# @Author: dongguangwen
# @Date  : 2025-03-23 20:51
from typing import Union

from langchain_openai import ChatOpenAI


model = ChatOpenAI(model="Pro-128K", api_key="dzCDkUHufCZEsuWlfVQd:",
                   base_url="https://spark-api-open.xf-yun.com/v1")


# Pydantic
class Joke(BaseModel):
    """Joke to tell user."""

    setup: str = Field(description="The setup of the joke")
    punchline: str = Field(description="The punchline to the joke")
    rating: Optional[int] = Field(
        default=None, description="How funny the joke is, from 1 to 10"
    )


class ConversationalResponse(BaseModel):
    """Respond in a conversational manner. Be kind and helpful."""

    response: str = Field(description="A conversational response to the user's query")


class FinalResponse(BaseModel):
    final_output: Union[Joke, ConversationalResponse]


structured_llm = llm.with_structured_output(FinalResponse)

structured_llm.invoke("Tell me a joke about cats")
